﻿#include "InjectFramework.h"
#include "FrameworkDataCenter.h"
#include "../common/pch.h"
#include "../interface/YunDunDefenderDlg.h"
#include <cjson/cJSON.h>
bool g_stop_process = false;

void database_process(std::shared_ptr<ProcessFramework> pf) {

    for (int i = 0; i < 100; i++)
    {
        // 处理操作
        Sleep(50);

        // 更新进度条
        pf->m_yddlg->SetScanDataBaseCTL(i + 1);
        if (i == 99)
        {
            i = 0;
        }
        // 检查是否需要退出
        if (g_stop_process)
        {
            break;
        }
    }
    pf->m_yddlg->SetScanDataBaseCTL(100);
}

void run_process_framework(std::shared_ptr<ProcessFramework> pf) {
    
    pf->Init();
    {
        std::lock_guard<std::mutex> lock(pf->process_mutex);
        g_stop_process = false;
        std::thread t(database_process, pf);
        t.detach();
        pf->m_dirqs->start_worker();
        pf->m_yddlg->SetScanWindowProcess(L"正在初始化数据库");
        pf->InitDataBase();
        g_stop_process = true;
    }
    pf->m_yddlg->SetScanWindowProcess(L"扫描进行中");
    pf->m_pathqs->start_worker();
    
}

void stop_process_framework(std::shared_ptr<ProcessFramework> pf) {
    pf->m_dirqs->stop_worker();
    pf->m_pathqs->stop_worker();
    pf->m_pathqs->clear_queue();
    pf->m_dirqs->clear_queue();
    pf->m_yddlg->SetScanWindowProcess(L"正在停止");
    {
        std::lock_guard<std::mutex> lock(pf->process_mutex);
        
        if (pf->m_engine)
        {
            cl_engine_free(pf->m_engine);
            pf->m_engine = nullptr;
        }
        
    }
    pf->m_yddlg->SetScanWindowProcess(L"引擎关闭");
    
   
    pf->m_yddlg->SetScanWindowProcess(L"扫描结束");
    pf->m_yddlg->SetScanProcessCTL(100);
    pf->m_yddlg->SetReuilt(pf->m_dwfilenumber, pf->m_allfienumber,pf->m_badFilenumber);
    Sleep(5000);
    pf->m_yddlg->SwitchTab(MENUTIME);
}

struct cb_context {
    const char* filename;
    unsigned long long virsize;
    char virhash[33];
    struct scan_cb_data* scandata;
};

void my_clcb_msg(enum cl_msg severity, const char* fullmsg, const char* msg, void* context)
{
    struct cb_context* c = (struct cb_context*)context;
    const char* filename = (c && c->filename) ? c->filename : "";

    UNUSEDPARAM(fullmsg);

    switch (severity) {
    case CL_MSG_ERROR:
        printf("^[LibClamAV] %s: %s", filename, msg);
        break;
    case CL_MSG_WARN:
        printf("~[LibClamAV] %s: %s", filename, msg);
        break;
    case CL_MSG_INFO_VERBOSE:
        printf("*[LibClamAV] %s: %s", filename, msg);
        break;
    default:
        printf("$[LibClamAV] %s: %s", filename, msg);
        break;
    }
}

ProcessFramework::ProcessFramework():m_pathqs(std::make_shared<PathQueueSet>()), m_dirqs(std::make_shared<DirQueueSet>())
{
    self_ptr = std::shared_ptr<ProcessFramework>(this, [](ProcessFramework*) {});
    m_pathqs->set_fkptr(self_ptr);
    m_dirqs->set_fkptr(self_ptr);
    int ret = cl_init(CL_INIT_DEFAULT);
    if (ret != CL_SUCCESS) {
        // 初始化失败，处理错误
    }
    //std::cout << "init defender ok" << std::endl;
    cl_set_clcb_msg(my_clcb_msg);
    m_ip = get_current_ip();
}


void ProcessFramework::SetYunDunDefenderDlg(std::shared_ptr<CYunDunDefenderDlg>& yddlg)
{
    m_yddlg = yddlg;
}

void ProcessFramework::Init()
{
    m_yddlg->SetScanWindowProcess(L"正在初始化");
    m_yddlg->SetScanWindowPath(L"");
    m_yddlg->SetScanProcessCTL(0);
    m_yddlg->SetScanDataBaseCTL(0);
    m_dwfilenumber = 0;
    m_allfienumber = 0;
    m_badFilenumber = 0;
    m_savepath = m_yddlg->GetSavePath().c_str();
    OpenOrCreateFile(m_savepath, outfile, std::ios::trunc);
    if (!outfile.is_open()) {
        std::cerr << "Error opening output file." << std::endl;
    }
    outfile.close();
}

void ProcessFramework::work()
{
    std::thread t(run_process_framework, self_ptr);
    t.detach();
}



int ProcessFramework::InitDataBase()
{
    int ret = 0;
    do
    {
        if (m_engine != 0)
        {
            ret = -1;
            break;
        }
        
        m_engine = cl_engine_new();
        if (!m_engine) {
            // 创建 engine 失败，处理错误
            ret = -1;
            break;
        }

        // 加载病毒库
        std::wstring str_p;
        wchar_t p[MAX_PATH + 1]{ 0 };
        GetModuleFileName(NULL, p, MAX_PATH);
        str_p = p;
        str_p = str_p.substr(0, str_p.rfind(L"\\"));
        str_p += L"\\database";

        m_signo = 0;

        ret = cl_load(WcharToUTF8(str_p.c_str()).c_str(), m_engine, &m_signo, CL_DB_STDOPT);

        if (ret != CL_SUCCESS) {
            // 加载病毒库失败，处理错误
            //std::cout << "read db  false, over!" << std::endl;
            break;
        }

        ret = cl_engine_compile(m_engine);

        if (CL_SUCCESS != ret) {
            //std::cout << "!reload_th: Database initialization error: can't compile engine: " << cl_strerror(ret) << std::endl;
            
            break;
        }

        cl_scan_options scanoptions;
        scanoptions.general = CL_SCAN_GENERAL_ALLMATCHES | CL_SCAN_GENERAL_COLLECT_METADATA | CL_SCAN_GENERAL_HEURISTICS | CL_SCAN_GENERAL_HEURISTIC_PRECEDENCE | CL_SCAN_GENERAL_UNPRIVILEGED;
        scanoptions.parse = CL_SCAN_PARSE_ARCHIVE | CL_SCAN_PARSE_ELF | CL_SCAN_PARSE_PDF | CL_SCAN_PARSE_SWF | CL_SCAN_PARSE_HWP3
            | CL_SCAN_PARSE_XMLDOCS | CL_SCAN_PARSE_MAIL | CL_SCAN_PARSE_OLE2 | CL_SCAN_PARSE_HTML | CL_SCAN_PARSE_PE;
        scanoptions.heuristic = CL_SCAN_HEURISTIC_BROKEN | CL_SCAN_HEURISTIC_EXCEEDS_MAX | CL_SCAN_HEURISTIC_PHISHING_SSL_MISMATCH | CL_SCAN_HEURISTIC_PHISHING_CLOAK
            | CL_SCAN_HEURISTIC_MACROS | CL_SCAN_HEURISTIC_ENCRYPTED_ARCHIVE | CL_SCAN_HEURISTIC_ENCRYPTED_DOC | CL_SCAN_HEURISTIC_PARTITION_INTXN | CL_SCAN_HEURISTIC_STRUCTURED
            | CL_SCAN_HEURISTIC_STRUCTURED_SSN_NORMAL | CL_SCAN_HEURISTIC_STRUCTURED_SSN_STRIPPED | CL_SCAN_HEURISTIC_STRUCTURED_CC | CL_SCAN_HEURISTIC_BROKEN_MEDIA;
        scanoptions.mail = CL_SCAN_MAIL_PARTIAL_MESSAGE;
        scanoptions.dev = CL_SCAN_DEV_COLLECT_SHA | CL_SCAN_DEV_COLLECT_PERFORMANCE_INFO;
    } while (false);
    return ret;
}


int ProcessFramework::ScanFile(std::wstring strpath)
{
    int ret = 0;
    do
    {
        std::lock_guard<std::mutex> lock(process_mutex);
        if (!m_engine)
        {
            break;
        }
        auto temppath = WcharToUTF8(strpath.c_str());
        const char* virname = nullptr;
        unsigned long int scanned = 0;
        struct cb_context context;
        context.filename = temppath.c_str();
        context.virsize = 0;
        auto ret = cl_scanfile_callback(temppath.c_str(), &virname, &scanned, m_engine, &m_scanoptions, &context);
        if (ret == CL_CLEAN)
        {
            //std::cout << "file is ok " << std::endl;
            //WriteData("unknow", scanned, temppath);
        }
        else if (ret == CL_VIRUS && (virname == NULL))
        {
            m_yddlg->AddScanColor(scanned);

            if (m_badFilenumber == 0)
            {
                m_yddlg->SetScanWindowProcess(L"扫描进行中(发现威胁！)");
            }
            
            m_badFilenumber++;
            // 文件被感染
            //std::cout << "file is bad ,virname unknow! " << std::endl;
            WriteData("unknow", scanned, temppath);
            
        }
        else if (ret == CL_VIRUS)
        {
            m_yddlg->AddScanColor(scanned);
            if (m_badFilenumber == 0)
            {
                m_yddlg->SetScanWindowProcess(L"扫描进行中(发现威胁！)");
            }
            m_badFilenumber++;
            //std::cout << "file is bad ,virname:  " << virname << " scanned: " << scanned << std::endl;
            WriteData(virname, scanned, temppath);
        }
    } while (false);
    return ret;
}

void ProcessFramework::AddDir(std::wstring strdir)
{
    m_dirqs->insert_data(strdir, strdir.size());
}

void ProcessFramework::AddPath(std::wstring strpath)
{
    if (strpath != L"end")
    {
        m_allfienumber++;
    }
    m_pathqs->insert_data(strpath, strpath.size());
}

void ProcessFramework::stop_work()
{
    std::thread t(stop_process_framework, self_ptr);
    t.detach();
}

void ProcessFramework::pause_work()
{
    m_dirqs->pause_worker();
    m_pathqs->pause_worker();
}

void ProcessFramework::resume_work()
{
    m_dirqs->resume_worker();
    m_pathqs->resume_worker();
}

std::string extractSubstring(const std::string& input) {
    std::size_t firstDot = input.find('.');
    std::size_t secondDot = input.find('.', firstDot + 1);

    if (firstDot != std::string::npos && secondDot != std::string::npos) {
        return input.substr(firstDot + 1, secondDot - firstDot - 1);
    }
    else {
        return "";
    }
}

void ProcessFramework::WriteData(std::string virname, unsigned long int scanned, std::string filename)
{
    cJSON* root = cJSON_CreateObject(); // 创建一个 JSON 对象

    cJSON_AddStringToObject(root, "IP", m_ip.c_str()); // 添加一个数字字段
    cJSON_AddStringToObject(root, "virusName", virname.c_str()); // 添加一个字符串字段
    std::string virusType;
    if (virname == "unknow")
    {
        virusType = "unknow";
    }
    else
    {
        virusType = extractSubstring(virname);
    }
    cJSON_AddStringToObject(root, "virusType", virusType.c_str()); // 添加一个数字字段
    cJSON_AddStringToObject(root, "virusPath", filename.c_str()); // 添加一个字符串字段
    std::string md5 = calculate_md5(filename);
    
    cJSON_AddStringToObject(root, "MD5", md5.c_str()); // 添加一个数字字段
    cJSON_AddStringToObject(root, "virusAbility", ""); // 添加一个数字字段
    cJSON_AddNumberToObject(root, "harmLevel", scanned); // 添加一个数字字段
    auto cur_time = get_current_utc_timestamp_ms();
    cJSON_AddNumberToObject(root, "issueDate", cur_time); // 添加一个数字字段
    cJSON_AddStringToObject(root, "details", ""); // 添加一个数字字段
    char* json_str = cJSON_PrintUnformatted(root); // 将 JSON 对象转换为未格式化的字符串
    outfile.open(m_savepath, std::ios::app);
    if (!outfile.is_open()) {
        std::cerr << "Error opening output file." << std::endl;
    }
    outfile << json_str << std::endl;
    outfile.close();
    cJSON_Delete(root); // 删除 JSON 对象，释放内存
    cJSON_free(json_str); // 释放 cJSON_PrintUnformatted 分配的内存
}